#include "AES.h"
#include <crypto++/aes.h>
#include <crypto++/modes.h>
#include <crypto++/filters.h>
#include "Base64.h"

namespace etu {
namespace common {

AES::AES(const std::string& source) {
    this->data = source;
}

AES::AES(const AES& orig) {
    this->data = orig.data;
}

AES::~AES() {
    //
}

std::string AES::encode(unsigned char* buffer, int length) {
    byte key[CryptoPP::AES::DEFAULT_KEYLENGTH];
    byte iv[CryptoPP::AES::BLOCKSIZE];

    memset(key, 0x00, CryptoPP::AES::DEFAULT_KEYLENGTH);
    memset(iv, 0x00, CryptoPP::AES::BLOCKSIZE);

    std::string ret;

    CryptoPP::AES::Encryption aesEncryption(key, CryptoPP::AES::DEFAULT_KEYLENGTH);
    CryptoPP::CBC_Mode_ExternalCipher::Encryption cbcEncryption(aesEncryption, iv);
    CryptoPP::StreamTransformationFilter stfEncryptor(cbcEncryption,
           new CryptoPP::StringSink(ret));

    stfEncryptor.Put(reinterpret_cast<const unsigned char*> ( buffer ),
            length + 1);
    stfEncryptor.MessageEnd();

    return Base64::encode(ret);
}

std::string AES::encodeString(const std::string& src) {
    byte key[CryptoPP::AES::DEFAULT_KEYLENGTH];
    byte iv[CryptoPP::AES::BLOCKSIZE];

    memset(key, 0x00, CryptoPP::AES::DEFAULT_KEYLENGTH);
    memset(iv, 0x00, CryptoPP::AES::BLOCKSIZE);

    std::string ret;

    CryptoPP::AES::Encryption aesEncryption(key, CryptoPP::AES::DEFAULT_KEYLENGTH);
    CryptoPP::CBC_Mode_ExternalCipher::Encryption cbcEncryption(aesEncryption, iv);
    CryptoPP::StreamTransformationFilter stfEncryptor(cbcEncryption,
           new CryptoPP::StringSink(ret));

    stfEncryptor.Put(reinterpret_cast<const unsigned char*> ( src.c_str() ),
            src.length() + 1);
    stfEncryptor.MessageEnd();

    return Base64::encode(ret);
}

void AES::decode(const std::string& encoded, unsigned char* buffer, int& length) {
    byte key[CryptoPP::AES::DEFAULT_KEYLENGTH];
    byte iv[CryptoPP::AES::BLOCKSIZE];

    memset(key, 0x00, CryptoPP::AES::DEFAULT_KEYLENGTH);
    memset(iv, 0x00, CryptoPP::AES::BLOCKSIZE);

    std::string src = Base64::decode(encoded);
    std::string ret;

    CryptoPP::AES::Decryption aesDecryption(key, CryptoPP::AES::DEFAULT_KEYLENGTH);
    CryptoPP::CBC_Mode_ExternalCipher::Decryption cbcDecryption( aesDecryption, iv );

    CryptoPP::StreamTransformationFilter stfDecryptor(cbcDecryption,
            new CryptoPP::StringSink(ret));
    stfDecryptor.Put( reinterpret_cast<const unsigned char*>(
        src.c_str() ), src.size());
    stfDecryptor.MessageEnd();

    length = src.length();
    buffer = (unsigned char *) src.c_str();
}

std::string AES::decodeString(const std::string& encoded) {
    byte key[CryptoPP::AES::DEFAULT_KEYLENGTH];
    byte iv[CryptoPP::AES::BLOCKSIZE];

    memset(key, 0x00, CryptoPP::AES::DEFAULT_KEYLENGTH);
    memset(iv, 0x00, CryptoPP::AES::BLOCKSIZE);

    std::string src = Base64::decode(encoded);
    std::string ret;

    CryptoPP::AES::Decryption aesDecryption(key, CryptoPP::AES::DEFAULT_KEYLENGTH);
    CryptoPP::CBC_Mode_ExternalCipher::Decryption cbcDecryption( aesDecryption, iv );

    CryptoPP::StreamTransformationFilter stfDecryptor(cbcDecryption,
            new CryptoPP::StringSink(ret));
    stfDecryptor.Put( reinterpret_cast<const unsigned char*>(
        src.c_str() ), src.size());
    stfDecryptor.MessageEnd();

    return ret;
}

std::ostream& AES::operator()(std::ostream& out) const {
    out << AES::encodeString(this->data);
    return out;
}

std::ostream& operator<<(std::ostream& out, AES aes) {
    return aes(out);
}

}
}
